#include "kernel_operator.h"
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;                                     // tensor num for each queue

//int
template<typename typeT> class KernelTruncDC {
public:
    __aicore__ inline KernelTruncDC() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, 
                                uint32_t CoreDataNum,
                                uint32_t finalTileNum,
                                uint32_t tileDataNum,
                                uint32_t TailDataNum) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->CoreDataNum = CoreDataNum;
        this->finalTileNum = finalTileNum;
        this->tileDataNum = tileDataNum;
        this->TailDataNum = TailDataNum;


        xGm.SetGlobalBuffer((__gm__ typeT*)x, this->CoreDataNum);
        yGm.SetGlobalBuffer((__gm__ typeT*)y, this->CoreDataNum);

        // pipe alloc memory to queue, the unit is Bytes
        pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileDataNum * sizeof(typeT));
        // pipe.InitBuffer(outQueueY, BUFFER_NUM, this->tilelength * sizeof(typeT));
        // pipe.InitBuffer(tmp1, this->tilelength * sizeof(float));
        // pipe.InitBuffer(tmp2, this->tilelength * sizeof(float));
        // pipe.InitBuffer(tmp3, this->tilelength * sizeof(int8_t) / 8);
    }
    __aicore__ inline void Process() {
        for (int32_t i = 0; i < this->finalTileNum - 1; i++) {
            CopyIn(i, this->tileDataNum);
            CopyOut(i, this->tileDataNum);
        }
        CopyIn(this->finalTileNum - 1, this->TailDataNum);
        CopyOut(this->finalTileNum - 1, this->TailDataNum);        
    }

private:
    __aicore__ inline void CopyIn(int32_t progress, uint32_t length) {
        LocalTensor<typeT> xLocal = inQueueX.AllocTensor<typeT>();
        DataCopy(xLocal, xGm[progress * this->tileDataNum], length);   
        inQueueX.EnQue(xLocal);
    }

    __aicore__ inline void CopyOut(int32_t progress, uint32_t length) {
        LocalTensor<typeT> xLocal = inQueueX.DeQue<typeT>();
        DataCopy(yGm[progress * this->tileDataNum], xLocal, length);  
        inQueueX.FreeTensor(xLocal);
    }

private:
    TPipe pipe;
    //TBuf<QuePosition::VECCALC> tmp1, tmp2, tmp3;
    TQueBind<TPosition::VECIN, TPosition::VECOUT, BUFFER_NUM> inQueueX;
    //TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    GlobalTensor<typeT> xGm;
    GlobalTensor<typeT> yGm;
    uint32_t CoreDataNum;
    uint32_t finalTileNum;
    uint32_t tileDataNum;
    uint32_t TailDataNum;

};

//bf16
template<typename typeT> class KernelTruncBF {
public:
    __aicore__ inline KernelTruncBF() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, 
                                uint32_t CoreDataNum,
                                uint32_t finalTileNum,
                                uint32_t tileDataNum,
                                uint32_t TailDataNum) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->CoreDataNum = CoreDataNum;
        this->finalTileNum = finalTileNum;
        this->tileDataNum = tileDataNum;
        this->TailDataNum = TailDataNum;


        xGm.SetGlobalBuffer((__gm__ typeT*)x, this->CoreDataNum);
        yGm.SetGlobalBuffer((__gm__ typeT*)y, this->CoreDataNum);

        // pipe alloc memory to queue, the unit is Bytes
        pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileDataNum * sizeof(typeT));
        pipe.InitBuffer(outQueueY, BUFFER_NUM, this->tileDataNum * sizeof(typeT));
        pipe.InitBuffer(tmp1, this->tileDataNum * sizeof(float));
        pipe.InitBuffer(tmp2, this->tileDataNum * sizeof(int64_t));
        // pipe.InitBuffer(tmp3, this->tilelength * sizeof(int8_t) / 8);
    }
    __aicore__ inline void Process() {
        for (int32_t i = 0; i < this->finalTileNum - 1; i++) {
            CopyIn(i, this->tileDataNum);
            Compute(this->tileDataNum);
            CopyOut(i, this->tileDataNum);
        }
        CopyIn(this->finalTileNum - 1, this->TailDataNum);
        Compute(this->TailDataNum);
        CopyOut(this->finalTileNum - 1, this->TailDataNum);        
    }

private:
    __aicore__ inline void CopyIn(int32_t progress, uint32_t length) {
        LocalTensor<typeT> xLocal = inQueueX.AllocTensor<typeT>();
        DataCopy(xLocal, xGm[progress * this->tileDataNum], length);   
        inQueueX.EnQue(xLocal);
    }

    __aicore__ inline void CopyOut(int32_t progress, uint32_t length) {
        LocalTensor<typeT> yLocal = outQueueY.DeQue<typeT>();
        DataCopy(yGm[progress * this->tileDataNum], yLocal, length);  
        outQueueY.FreeTensor(yLocal);
    }
    __aicore__ inline void Compute(uint32_t length) {
        LocalTensor<typeT> xLocal = inQueueX.DeQue<typeT>();
        LocalTensor<typeT> yLocal = outQueueY.AllocTensor<typeT>();
        LocalTensor<float> tmp1Local = tmp1.Get<float>();
        LocalTensor<int64_t> tmp2Local = tmp2.Get<int64_t>();
        if constexpr (std::is_same_v<typeT, __bf16>) {
            Cast(tmp1Local, xLocal, RoundMode::CAST_NONE, length);
            Cast(tmp2Local, tmp1Local, RoundMode::CAST_TRUNC, length);
            Cast(tmp1Local, tmp2Local, RoundMode::CAST_RINT, length);
            Cast(yLocal, tmp1Local, RoundMode::CAST_RINT, length);
        }
        outQueueY.EnQue<typeT>(yLocal);
        inQueueX.FreeTensor(xLocal);
    }

private:
    TPipe pipe;
    TBuf<QuePosition::VECCALC> tmp1, tmp2;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    GlobalTensor<typeT> xGm;
    GlobalTensor<typeT> yGm;
    uint32_t CoreDataNum;
    uint32_t finalTileNum;
    uint32_t tileDataNum;
    uint32_t TailDataNum;

};

//half
template<typename typeT> class KernelTruncHALF {
public:
    __aicore__ inline KernelTruncHALF() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, 
                                uint32_t CoreDataNum,
                                uint32_t finalTileNum,
                                uint32_t tileDataNum,
                                uint32_t TailDataNum) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->CoreDataNum = CoreDataNum;
        this->finalTileNum = finalTileNum;
        this->tileDataNum = tileDataNum;
        this->TailDataNum = TailDataNum;

        xGm.SetGlobalBuffer((__gm__ typeT*)x, this->CoreDataNum);
        yGm.SetGlobalBuffer((__gm__ typeT*)y, this->CoreDataNum);

        // pipe alloc memory to queue, the unit is Bytes
        pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileDataNum * sizeof(typeT));
        pipe.InitBuffer(outQueueY, BUFFER_NUM, this->tileDataNum * sizeof(typeT));
        pipe.InitBuffer(tmp1, this->tileDataNum * sizeof(float));
        pipe.InitBuffer(tmp2, this->tileDataNum * sizeof(int32_t));
        // pipe.InitBuffer(tmp3, this->tilelength * sizeof(int8_t) / 8);
    }
    __aicore__ inline void Process() {
        for (int32_t i = 0; i < this->finalTileNum - 1; i++) {
            CopyIn(i, this->tileDataNum);
            Compute(this->tileDataNum);
            CopyOut(i, this->tileDataNum);
        }
        CopyIn(this->finalTileNum - 1, this->TailDataNum);
        Compute(this->TailDataNum);
        CopyOut(this->finalTileNum - 1, this->TailDataNum);        
    }

private:
    __aicore__ inline void CopyIn(int32_t progress, uint32_t length) {
        LocalTensor<typeT> xLocal = inQueueX.AllocTensor<typeT>();
        DataCopy(xLocal, xGm[progress * this->tileDataNum], length);   
        inQueueX.EnQue(xLocal);
    }

    __aicore__ inline void CopyOut(int32_t progress, uint32_t length) {
        LocalTensor<typeT> yLocal = outQueueY.DeQue<typeT>();
        DataCopy(yGm[progress * this->tileDataNum], yLocal, length);  
        outQueueY.FreeTensor(yLocal);
    }
    __aicore__ inline void Compute(uint32_t length) {
        LocalTensor<typeT> xLocal = inQueueX.DeQue<typeT>();
        LocalTensor<typeT> yLocal = outQueueY.AllocTensor<typeT>();
        LocalTensor<float> tmp1Local = tmp1.Get<float>();
        LocalTensor<int32_t> tmp2Local = tmp2.Get<int32_t>();
        if constexpr (std::is_same_v<typeT, half>) {
            Cast(tmp2Local, xLocal, RoundMode::CAST_TRUNC, length);
            Cast(tmp1Local, tmp2Local, RoundMode::CAST_RINT, length);
            Cast(yLocal, tmp1Local, RoundMode::CAST_RINT, length);
        }
        outQueueY.EnQue<typeT>(yLocal);
        inQueueX.FreeTensor(xLocal);
    }

private:
    TPipe pipe;
    TBuf<QuePosition::VECCALC> tmp1, tmp2;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    GlobalTensor<typeT> xGm;
    GlobalTensor<typeT> yGm;
    uint32_t CoreDataNum;
    uint32_t finalTileNum;
    uint32_t tileDataNum;
    uint32_t TailDataNum;

};


//FLOAT32
template<typename typeT> class KernelTruncFLOAT {
public:
    __aicore__ inline KernelTruncFLOAT() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, 
                                uint32_t CoreDataNum,
                                uint32_t finalTileNum,
                                uint32_t tileDataNum,
                                uint32_t TailDataNum) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->CoreDataNum = CoreDataNum;
        this->finalTileNum = finalTileNum;
        this->tileDataNum = tileDataNum;
        this->TailDataNum = TailDataNum;

        xGm.SetGlobalBuffer((__gm__ typeT*)x, this->CoreDataNum);
        yGm.SetGlobalBuffer((__gm__ typeT*)y, this->CoreDataNum);

        // pipe alloc memory to queue, the unit is Bytes
        pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileDataNum * sizeof(typeT));
        pipe.InitBuffer(outQueueY, BUFFER_NUM, this->tileDataNum * sizeof(typeT));
        // pipe.InitBuffer(tmp1, this->tileDataNum * sizeof(float));
        pipe.InitBuffer(tmp2, this->tileDataNum * sizeof(int64_t));
        // pipe.InitBuffer(tmp3, this->tilelength * sizeof(int8_t) / 8);
    }
    __aicore__ inline void Process() {
        for (int32_t i = 0; i < this->finalTileNum - 1; i++) {
            CopyIn(i, this->tileDataNum);
            Compute(this->tileDataNum);
            CopyOut(i, this->tileDataNum);
        }
        CopyIn(this->finalTileNum - 1, this->TailDataNum);
        Compute(this->TailDataNum);
        CopyOut(this->finalTileNum - 1, this->TailDataNum);        
    }

private:
    __aicore__ inline void CopyIn(int32_t progress, uint32_t length) {
        LocalTensor<typeT> xLocal = inQueueX.AllocTensor<typeT>();
        DataCopy(xLocal, xGm[progress * this->tileDataNum], length);   
        inQueueX.EnQue(xLocal);
    }

    __aicore__ inline void CopyOut(int32_t progress, uint32_t length) {
        LocalTensor<typeT> yLocal = outQueueY.DeQue<typeT>();
        DataCopy(yGm[progress * this->tileDataNum], yLocal, length);  
        outQueueY.FreeTensor(yLocal);
    }
    __aicore__ inline void Compute(uint32_t length) {
        LocalTensor<typeT> xLocal = inQueueX.DeQue<typeT>();
        LocalTensor<typeT> yLocal = outQueueY.AllocTensor<typeT>();
        // LocalTensor<float> tmp1Local = tmp1.Get<float>();
        LocalTensor<int64_t> tmp2Local = tmp2.Get<int64_t>();
        if constexpr (std::is_same_v<typeT, float>) {
            Cast(tmp2Local, xLocal, RoundMode::CAST_TRUNC, length);
            Cast(yLocal, tmp2Local, RoundMode::CAST_RINT, length);
        }
        outQueueY.EnQue<typeT>(yLocal);
        inQueueX.FreeTensor(xLocal);
    }

private:
    TPipe pipe;
    TBuf<QuePosition::VECCALC> tmp2;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    GlobalTensor<typeT> xGm;
    GlobalTensor<typeT> yGm;
    uint32_t CoreDataNum;
    uint32_t finalTileNum;
    uint32_t tileDataNum;
    uint32_t TailDataNum;
};

extern "C" __global__ __aicore__ void trunc(GM_ADDR input_x, GM_ADDR output_y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    if (TILING_KEY_IS(200)) {
        KernelTruncDC<DTYPE_INPUT_X> op;
        op.Init(input_x, output_y, 
            tiling_data.CoreDataNum, tiling_data.finalTileNum, tiling_data.tileDataNum, tiling_data.TailDataNum);  
        op.Process();
    } else if (TILING_KEY_IS(100)) {
        KernelTruncBF<DTYPE_INPUT_X> op;
        op.Init(input_x, output_y, 
            tiling_data.CoreDataNum, tiling_data.finalTileNum, tiling_data.tileDataNum, tiling_data.TailDataNum);  
        op.Process();
    } else if (TILING_KEY_IS(300)) {
        KernelTruncHALF<DTYPE_INPUT_X> op;
        op.Init(input_x, output_y, 
            tiling_data.CoreDataNum, tiling_data.finalTileNum, tiling_data.tileDataNum, tiling_data.TailDataNum);  
        op.Process();        
    } else if (TILING_KEY_IS(400)) {
        KernelTruncFLOAT<DTYPE_INPUT_X> op;
        op.Init(input_x, output_y, 
            tiling_data.CoreDataNum, tiling_data.finalTileNum, tiling_data.tileDataNum, tiling_data.TailDataNum);  
        op.Process();
    }
}